# -*- coding: utf-8 -*-
"""
文件: src/agents/medgemma_agent.py

用途：
  - 以 MedGemma（或兼容的多模态/文本模型）为底座，构造双角色 Agent：
    * Generator：面向眼底阅片知识的问答生成
    * Evaluator：对生成的问答进行医学一致性与可用性校验
  - 支持三类生成任务：
    1) 疾病单点问答
    2) 两病关系/易混问答
    3) 诊断推演问答
说明：
  - 本版本仅支持**纯文本生成**（不输入图像）。
  - 自动识别 4B / 27B 型号并选择最优 device_map 策略。
"""

import os
import re
import json
import random
from typing import Dict, Any, List, Optional
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor, AutoTokenizer

# ========== 默认值 ==========
DEFAULT_MODEL_ID = "google/medgemma-4b-it"
DEFAULT_DTYPE = torch.bfloat16
DEFAULT_TEMPERATURE = 0.9
DEFAULT_TOP_P = 0.9
DEFAULT_MAX_NEW_TOKENS = 768

# ============================================================
# 通用系统提示
# ============================================================
SYSTEM_PROMPT = (
    "你是一名资深眼科医生与阅片专家，擅长通过眼底影像判断疾病特征。"
    "在回答问题时，请关注视觉表现、病灶形态、病程变化及常见误判点。"
    "保持语言准确、客观、专业。"
)

# ========= Evaluator 鲁棒解析 =========

def _strip_code_fence(text: str) -> str:
    if not isinstance(text, str):
        return ""
    s = text.strip()
    if s.startswith("```"):
        s = re.sub(r"^```[a-zA-Z0-9]*\s*", "", s)
        s = re.sub(r"\s*```$", "", s)
    return s.strip()

def _safe_json_load(s: str):
    try:
        return json.loads(s)
    except Exception:
        return None

def _coerce_float(x) -> Optional[float]:
    if x is None:
        return None
    if isinstance(x, (int, float)):
        return float(x)
    if isinstance(x, str):
        xs = x.strip().strip('"').strip("'")
        m = re.search(r"[-+]?\d+(\.\d+)?", xs)
        if m:
            try:
                return float(m.group(0))
            except Exception:
                return None
    return None

def _norm_score_to_0_1(score: Optional[float]) -> Optional[float]:
    if score is None:
        return None
    # 支持 0~1 或 0~100
    if 0.0 <= score <= 1.0:
        return score
    if 1.0 < score <= 100.0:
        return score / 100.0
    return None

def _norm_decision(s: Optional[str]) -> Optional[str]:
    if not s:
        return None
    t = s.strip().lower()
    # 常见同义词
    if t in ("accept", "accepted", "是", "通过", "同意", "接受", "true", "yes", "ok", "合格"):
        return "accept"
    if t in ("reject", "rejected", "否", "不通过", "拒绝", "false", "no", "不合格"):
        return "reject"
    if t in ("uncertain", "unsure", "不确定", "存疑", "待定"):
        return "uncertain"
    return None

def _kv_capture(block: str, keys: List[str]) -> Optional[str]:
    """
    从文本里按 `key: value` 或 `key：value` 抓取值；返回捕获到的 value（去掉首尾空白）
    """
    if not isinstance(block, str):
        return None
    for k in keys:
        # 允许中文/英文冒号，允许运行到行尾或遇到下一行
        m = re.search(rf"{k}\s*[:：]\s*(.+)", block, flags=re.I)
        if m:
            return m.group(1).strip()
    return None

def parse_eval_loose(text: str) -> Dict[str, Any]:
    """
    解析 Evaluator 输出（尽量）。输出字段：
      - decision: accept/reject/uncertain 或 None
      - score: 0~1 浮点 或 None
      - rationale: str 或 ""
      - parsed_from: "json" | "kv" | "fallback"
      - 兼容字段：is_correct, confidence
    """
    res = {
        "decision": None,
        "score": None,
        "rationale": "",
        "parsed_from": "fallback",
        # 兼容旧接口：
        "is_correct": "uncertain",
        "confidence": 0.0,
    }

    if not isinstance(text, str) or not text.strip():
        return res

    s = _strip_code_fence(text)

    # 1) 先试 JSON
    js = None
    # 1.1 如果包含花括号块，只取第一段 {...}
    brace = re.search(r"\{[\s\S]*\}", s)
    if brace:
        js = _safe_json_load(brace.group(0))
    else:
        js = _safe_json_load(s)
    if isinstance(js, dict):
        score = _coerce_float(js.get("score") or js.get("confidence"))
        score = _norm_score_to_0_1(score)
        decision = _norm_decision(js.get("decision") or js.get("verdict"))
        rationale = js.get("rationale") or js.get("reason") or js.get("explanation") or ""
        res["score"] = score
        res["decision"] = decision
        res["rationale"] = str(rationale).strip()
        res["parsed_from"] = "json"
    else:
        # 2) 键值行解析
        score_raw = _kv_capture(s, ["打分", "评分", "score", "分数", "confidence", "置信度"])
        decision_raw = _kv_capture(s, ["是否接受", "是否通过", "decision", "verdict", "判定", "结论"])
        rationale_raw = _kv_capture(s, ["理由", "依据", "说明", "原因", "rationale", "reason", "explanation"])

        score = _norm_score_to_0_1(_coerce_float(score_raw))
        decision = _norm_decision(decision_raw)

        # 如果找到了“打分/是否接受”但值不规范，依然把原始片段塞进 rationale 里，避免信息丢失
        if (score is None) and score_raw:
            rationale_raw = f"【原始打分不可解析】{score_raw}\n" + (rationale_raw or "")
        if (decision is None) and decision_raw:
            rationale_raw = f"【原始判定不可解析】{decision_raw}\n" + (rationale_raw or "")

        res["score"] = score
        res["decision"] = decision
        res["rationale"] = (rationale_raw or "").strip()
        res["parsed_from"] = "kv" if (score_raw or decision_raw or rationale_raw) else "fallback"

        # 3) 仍没抓到 rationale，就截取前 400 字当说明
        if not res["rationale"]:
            res["rationale"] = s[:400].strip()

    # ---- 兼容旧接口字段 ----
    # 优先用 score 推出 confidence
    conf = res["score"]
    if conf is None:
        # 从文本里再粗暴找一个 0~1 数字当置信度（可选）
        mm = re.search(r"(?<!\d)(0\.\d{1,3}|1\.0)\b", s)
        conf = float(mm.group(1)) if mm else 0.0

    decision = res["decision"]
    if decision == "accept":
        is_correct = True
    elif decision == "reject":
        is_correct = False
    else:
        is_correct = "uncertain"

    res["is_correct"] = is_correct
    res["confidence"] = float(conf)

    return res

# ============================================================
# Prompt 模板（强化视觉表达）
# ============================================================
GENERATOR_HEADER = '''
你是一名眼底阅片专家。你的任务是围绕眼底影像的**可见视觉线索**生成高质量中文问答，帮助医生进行诊断与鉴别。要求：
- 严格使用**纯中文**术语（必要时可用“（英文缩写）”补充，如“视网膜神经纤维层（RNFL）”），**不得出现纯英文词条**；
- **不得使用任何 Markdown 代码块或 ```json 围栏**；
- 注重“看到什么→提示什么→还需如何确认”的链式表达，避免绝对化结论；
- 优先描述：结构（位置/形态/边缘）、颜色/反光、血管走向、量化线索（如“与对侧/基线相比杯盘比增大”）、伴随体征；
- 确保每条问答**不重复**、**信息独立**、**符合临床阅片逻辑**。
'''

GEN_TASK_SINGLE_DISEASE = '''
【任务1：医学知识问答】
目标疾病：{disease_cn}

请生成 {n} 组问答，每组从不同角度介绍该疾病的核心知识点，偏重“图像可见线索”。必须满足：
1) 问题像医学生课堂提问，自然、具体，长度 10~40 字；
2) 回答专业、清晰，长度 60~200 字，包含：定义/流行病学要点（简）、典型眼底视觉特征、如何据此做出判断（含“与对侧/基线比较”等谨慎措辞）；
3) 如该病**确有**分期/分型的国际通行标准（AAO/ICO 等），可在其中一条中**简述原则**（避免罗列全部条款）；
4) 仅输出**严格 JSON 数组**，每个元素含 "question" 和 "answer" 两个键；
5) 禁止输出任何解释、编号、注释、代码块或 ```json。

【输出示例】
[
  {
    "question": "青光眼在阅片时应首先关注哪些结构性线索？",
    "answer": "重点观察杯盘比相较对侧或基线是否增大、视神经缘部是否变薄（尤以上下象限）、是否存在楔形/弧形的视网膜神经纤维层缺损（RNFL），以及血管是否呈鼻向偏移。上述线索的组合可提示青光眼性视神经损害，需结合眼压、对侧眼及视野进一步评估。"
  }
]
【硬性要求】仅输出以 "[" 开始、以 "]" 结束的 JSON 数组，不要出现 ```json。
'''

GEN_TASK_RELATION_PAIR = '''【任务2：两病鉴别】
疾病A：{a_cn}
疾病B：{b_cn}

请生成 {n} 组问答，围绕医生在**眼底图像**中如何区分这两种疾病。必须满足：
1) 每个问题以**视觉差异**为切入点（颜色/反光、形态/边缘、病灶分布、血管改变、是否累及黄斑/视盘等）；
2) 回答需指出关键差异并给出**诊断意义**，避免绝对化，必要时提示“需结合病史/检查”；
3) 仅输出**严格 JSON 数组**，每个元素含 "question" 和 "answer" 两个键；
4) 禁止输出解释文字、英文夹杂、Markdown 代码块或 ```json。

【输出示例】
[
  {
    "question": "糖尿病视网膜病变与高血压视网膜病变在病灶类型与分布上如何区分？",
    "answer": "糖尿病视网膜病变常见微动脉瘤、硬性渗出、黄斑区受累及新生血管；高血压视网膜病变更常见动脉反光增强、动静脉交叉压迫、火焰状出血与棉絮斑，分布更弥散。若伴视盘水肿，需结合血压与全身状况进一步判断。"
  }
]
【硬性要求】仅输出以 "[" 开始、以 "]" 结束的 JSON 数组，不要出现 ```json。
'''
GEN_TASK_REASONING ='''【任务3：阅片推演】
疾病：{disease_cn}

请生成 {n} 组问答，模拟临床阅片推理。必须满足：
1) 问题用一句话描述**观察到的具体视觉体征组合**（如“杯盘比较对侧显著增大+视盘缘部变薄+RNFL 楔形缺损”），并询问“这提示什么/接下来该如何判断”；
2) 回答直接给出**可能诊断/阶段倾向**与**进一步核实建议**（如眼压/对侧比较/视野/OCT）；
3) 仅输出**严格 JSON 数组**，每个元素含 "question" 和 "answer" 两个键；禁止 tags/observation 等键；
4) 禁止输出任何解释或 Markdown 代码块/```json。

【输出示例】
[
  {
    "question": "若看到杯盘比较对侧明显增大，视神经缘部上下象限变薄并伴 RNFL 楔形缺损，这提示什么？",
    "answer": "上述组合提示青光眼性视神经损害的可能性较高。建议结合眼压、对侧眼杯盘比、视野与 OCT 进一步确认，并动态随访评估进展。"
  }
]
【硬性要求】仅输出以 "[" 开始、以 "]" 结束的 JSON 数组，不要出现 ```json。
'''

EVALUATOR_HEADER = '''
你是一名资深眼科医生与阅片审校专家。请审阅一对中文问答(Q/A)，只围绕“眼底**视觉线索**正确性、相关性、表达清晰度”作出判断。

尽量输出**一个 JSON 对象**（推荐键）：
{
  "score": 0.0~1.0 ,          # 评分，越高越好
  "decision": "accept/reject/uncertain",  # 是否接受
  "rationale": "简要理由"                 # 200字内
}
禁止输出 Markdown 代码块或 ```json。
'''
EVAL_PROMPT_TEMPLATE = (
    EVALUATOR_HEADER + "\n\n"
    "前置信息：{context}\n"
    "Q: {question}\n"
    "A: {answer}\n"
    "请给出上述结构化结论。"
)


# ============================================================
# Agent 类
# ============================================================
class MedGemmaAgent:
    def __init__(self, model_id=DEFAULT_MODEL_ID, role="generator",
                 dtype=DEFAULT_DTYPE, temperature=DEFAULT_TEMPERATURE,
                 top_p=DEFAULT_TOP_P, max_new_tokens=DEFAULT_MAX_NEW_TOKENS):
        assert role in ("generator", "evaluator")
        self.role = role
        self.dtype = dtype
        self.temperature = temperature
        self.top_p = top_p
        self.max_new_tokens = max_new_tokens

        # ✅ 智能判断加载方式
        is_large = any(k in model_id.lower() for k in ["27b", "30b", "70b"])
        if is_large:
            print(f"[INFO] Loading large text-only model '{model_id}' (device_map='auto')")
            self.model = AutoModelForImageTextToText.from_pretrained(
                model_id, device_map="auto",
                torch_dtype=dtype, low_cpu_mem_usage=True)
        else:
            print(f"[INFO] Loading multimodal model '{model_id}' on cuda:0")
            self.model = AutoModelForImageTextToText.from_pretrained(
                model_id, torch_dtype=dtype,
                low_cpu_mem_usage=True).to("cuda:0")

        # ✅ Processor/Tokenizer
        try:
            self.processor = AutoProcessor.from_pretrained(model_id)
        except Exception:
            self.processor = AutoTokenizer.from_pretrained(model_id, use_fast=True)

        self._first_device = self._infer_first_device()
        print(f"[INFO] First active device: {self._first_device}")

    # ------- 设备工具 -------
    def _infer_first_device(self):
        if hasattr(self.model, "hf_device_map") and isinstance(self.model.hf_device_map, dict):
            for _, dev in self.model.hf_device_map.items():
                if isinstance(dev, str) and dev.startswith("cuda"):
                    return torch.device(dev)
        return next(self.model.parameters()).device

    def _move_inputs(self, inputs, device):
        moved = {}
        for k, v in inputs.items():
            if hasattr(v, "to"):
                moved[k] = v.to(device=device, dtype=self.dtype if torch.is_floating_point(v) else None)
            else:
                moved[k] = v
        return moved

    # ------- 通用推理 -------
    def _chat(self, messages: List[Dict[str, Any]]) -> str:
        inputs = self.processor.apply_chat_template(
            messages, add_generation_prompt=True,
            tokenize=True, return_dict=True, return_tensors="pt"
        )
        inputs = self._move_inputs(inputs, self._first_device)

        input_len = inputs["input_ids"].shape[-1]
        with torch.inference_mode():
            gen = self.model.generate(
                **inputs,
                max_new_tokens=self.max_new_tokens,
                do_sample=True,
                temperature=self.temperature,
                top_p=self.top_p
            )
        decoded = self.processor.decode(gen[0][input_len:], skip_special_tokens=True)
        return decoded.strip()

    # ------- 三类生成 -------
    def generate_single_disease(self, disease_cn: str, n: int = 6) -> List[Dict[str, Any]]:
        task_prompt = GEN_TASK_SINGLE_DISEASE.replace("{disease_cn}", disease_cn).replace("{n}", str(n))
        messages = [
            {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
            {"role": "user", "content": [{"type": "text", "text": GENERATOR_HEADER + "\n" + task_prompt}]}
        ]
        raw = self._chat(messages)
        try:
            return json.loads(self._extract_json_array(raw))
        except Exception:
            
            print("paser json fail")
            print(raw)
            return []

    def generate_relation_pair(self, disease_a_cn: str, disease_b_cn: str, n: int = 6) -> List[Dict[str, Any]]:
        """
        任务2：两病鉴别问答生成
        """
        # ⚙️ 构造任务提示词（安全字符串替换）
        task_prompt = (
            GEN_TASK_RELATION_PAIR
            .replace("{a_cn}", disease_a_cn)
            .replace("{b_cn}", disease_b_cn)
            .replace("{n}", str(n))
        )

        messages = [
            {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
            {"role": "user", "content": [{"type": "text", "text": GENERATOR_HEADER + "\n" + task_prompt}]}
        ]

        raw = self._chat(messages)
        try:
            return json.loads(self._extract_json_array(raw))
        except Exception as e:
            print(f"[WARN] JSON parse failed in generate_relation_pair: {e}")
            print("Raw output snippet:\n", raw[:800])
            return []


    def generate_reasoning(self, disease_cn: str, n: int = 6) -> List[Dict[str, Any]]:
        """
        任务3：阅片推演问答生成
        """
        # ⚙️ 安全替换构造提示
        task_prompt = (
            GEN_TASK_REASONING
            .replace("{disease_cn}", disease_cn)
            .replace("{n}", str(n))
        )

        messages = [
            {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
            {"role": "user", "content": [{"type": "text", "text": GENERATOR_HEADER + "\n" + task_prompt}]}
        ]

        raw = self._chat(messages)
        try:
            return json.loads(self._extract_json_array(raw))
        except Exception as e:
            print(f"[WARN] JSON parse failed in generate_reasoning: {e}")
            print("Raw output snippet:\n", raw[:800])
            raise ValueError("JSON parse failed in generate_reasoning")


    # ------- 评估 -------
    def evaluate_qa(self, question: str, answer: str, context: Optional[str] = None) -> Dict[str,  Any]:
        """
        评估生成问答质量（鲁棒解析）：
          - 首选 JSON（score/decision/rationale）
          - 次选 键值行：打分/是否接受/理由
          - 兜底：从原文抢救出尽可能多的信息
        返回包含新老两套字段：
          - 新: decision, score, rationale, parsed_from
          - 旧: is_correct(True/False/"uncertain"), confidence(0~1)
        """
        assert self.role == "evaluator"

        header = EVALUATOR_HEADER
        prompt = (
        EVAL_PROMPT_TEMPLATE
        .replace("{context}", context or "无")
        .replace("{question}", question)
        .replace("{answer}", answer)
        )
        inputs_content= header + "\n\n" + prompt
        # print("===== EVAL PROMPT =====")
        # print(inputs_content)
        messages = [
            {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
            {"role": "user", "content": [{"type": "text", "text": inputs_content}]}
        ]
        raw = self._chat(messages)

        # 新的鲁棒解析
        parsed = parse_eval_loose(raw)

        # 也把原文带回，方便调试（不会影响你的打分逻辑）
        parsed["raw"] = raw
        return parsed



    # ------- 工具 -------
    @staticmethod
    def _extract_json_array(text: str) -> str:
        """
        Robust JSON array extractor that supports:
        - ```json fenced code blocks
        - Truncated or missing closing ```
        - Extra text/logs before or after
        - Trailing commas or partial last elements
        """

        # 1. 如果存在 ```json，则以其为起点
        fence_start = text.find("```json")
        if fence_start != -1:
            start = fence_start + len("```json")
            end = text.find("```", start + 1)
            if end == -1:
                # 无闭合 fence，读到文本末尾
                json_block = text[start:]
            else:
                json_block = text[start:end]
        else:
            # 无 ```json，尝试匹配第一个 JSON 数组
            match = re.search(r"\[.*\]", text, flags=re.S)
            if match:
                json_block = match.group(0)
            else:
                raise ValueError("No JSON array found in model output.")

        # 2. 清理多余空格
        s = json_block.strip()

        # 3. 去掉可能多余的说明性行
        s = re.sub(r"^[^\[]*\[", "[", s, flags=re.S)

        # 4. 修复尾部逗号/未闭合 JSON 对象
        s = re.sub(r",\s*]", "]", s)
        # 若最后一个对象没闭合，尝试闭合
        if s.count("{") > s.count("}"):
            s += "}" * (s.count("{") - s.count("}"))
        if not s.strip().endswith("]"):
            s += "]"

        # 5. 尝试 parse；若失败再宽松修正一次
        try:
            json.loads(s)
            return s
        except Exception:
            # 尝试去掉最后一个半截对象
            s2 = re.sub(r",\s*{\s*[^}\]]*$", "]", s)
            json.loads(s2)
            return s2

# ============================================================
# Demo 入口（测试）
# ============================================================
if __name__ == "__main__":
    # ============================================================
    # Demo: 基础生成 + 评估测试
    # ============================================================

    print("===== 🩺 MedGemma Agent 测试开始 =====")

    # 初始化生成器与评估器
    gen_agent = MedGemmaAgent(model_id="./pretrained/medgemma4b", role="generator")
    eva_agent = MedGemmaAgent(model_id="./pretrained/medgemma4b", role="evaluator")

    # ------------------------------------------------------------
    # ① 单病问答生成
    print("\n=== [任务1] 单病问答：青光眼 ===")
    result_single = gen_agent.generate_single_disease("青光眼", n=3)
    print(json.dumps(result_single, ensure_ascii=False, indent=2))

    # 评估第一条
    if result_single:
        qa = result_single[0]
        eval_res = eva_agent.evaluate_qa(qa["question"], qa["answer"], context="任务1-青光眼单病")
        print("\n--- [评估结果] ---")
        print(json.dumps(eval_res, ensure_ascii=False, indent=2))
    else:
        print("[WARN] 单病问答生成为空，跳过评估")

    print("-" * 60)

    # ------------------------------------------------------------
    # ② 两病关系问答
    print("\n=== [任务2] 两病关系问答：糖尿病视网膜病变 vs 高血压视网膜病变 ===")
    result_relation = gen_agent.generate_relation_pair("糖尿病视网膜病变", "高血压视网膜病变", n=2)
    print(json.dumps(result_relation, ensure_ascii=False, indent=2))

    # 评估第一条
    if result_relation:
        qa = result_relation[0]
        eval_res = eva_agent.evaluate_qa(qa["question"], qa["answer"], context="任务2-两病关系")
        print("\n--- [评估结果] ---")
        print(json.dumps(eval_res, ensure_ascii=False, indent=2))
    else:
        print("[WARN] 两病关系生成为空，跳过评估")

    print("-" * 60)

    # ------------------------------------------------------------
    # ③ 推理问答
    print("\n=== [任务3] 阅片推演：青光眼 ===")
    result_reasoning = gen_agent.generate_reasoning("青光眼", n=2)
    print(json.dumps(result_reasoning, ensure_ascii=False, indent=2))

    # 评估第一条
    if result_reasoning:
        qa = result_reasoning[0]
        eval_res = eva_agent.evaluate_qa(qa["question"], qa["answer"], context="任务3-推演问答")
        print("\n--- [评估结果] ---")
        print(json.dumps(eval_res, ensure_ascii=False, indent=2))
    else:
        print("[WARN] 推理问答生成为空，跳过评估")

    print("===== ✅ 测试完成 =====")
